// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/common/file/temp_directory_path.h>
#include <pollux/runner/local_runner.h>

namespace kumo::pollux::exec::test {

struct TableSpec {
  std::string name;
  RowTypePtr columns;
  int32_t rowsPerVector{10000};
  int32_t numVectorsPerFile{5};
  int32_t numFiles{5};

  /// Function  Applied to generated RowVectors for the table before writing.
  /// May be used to insert non-random data on top of the random datafrom
  /// HiveConnectorTestBase::makeVectors.
  std::function<void(const RowVectorPtr& vector)> customizeData;
};

/// Test helper class that manages a TestCase with a set of generated
/// tables and a HiveConnector that exposes the files and their
/// metadata. The lifetime the test data is the test case consisting
/// of multiple google unit test cases.
class LocalRunnerTestBase : public HiveConnectorTestBase {
 protected:
  static void SetUpTestCase() {
    HiveConnectorTestBase::SetUpTestCase();
    schemaExecutor_ = std::make_unique<melon::CPUThreadPoolExecutor>(4);
  }

  static void TearDownTestCase() {
    initialized_ = false;
    files_.reset();
    HiveConnectorTestBase::TearDownTestCase();
  }

  void SetUp() override;

  void ensureTestData();

  /// Re-creates the connector with kHiveConnectorId with a config
  /// that points to the temp directory created by 'this'. If the
  /// connector factory is wired to capture metadata then the metadata
  /// will be available through the connector.
  void setupConnector();

  /// Returns a split source factory that contains splits for the table scans in
  /// 'plan'. 'plan' should refer to testing tables created by 'this'.
  std::shared_ptr<runner::SimpleSplitSourceFactory>
  makeSimpleSplitSourceFactory(const runner::MultiFragmentPlanPtr& plan);

  void makeTables(
      std::vector<TableSpec> specs,
      std::shared_ptr<TempDirectoryPath>& directory);

  // Creates a QueryCtx with 'pool'. 'pool' must be a root pool.
  static std::shared_ptr<core::QueryCtx> makeQueryCtx(
      const std::string& queryId,
      memory::MemoryPool* pool);

  // Configs for creating QueryCtx.
  inline static std::unordered_map<std::string, std::string> config_;
  inline static std::unordered_map<std::string, std::string> hiveConfig_;

  // The specification of the test data. The data is created in ensureTestData()
  // called from each SetUp()(.
  inline static std::vector<TableSpec> testTables_;

  // The top level directory with the test data.
  inline static bool initialized_;
  inline static std::string testDataPath_;
  inline static std::string localFileFormat_{"dwrf"};
  inline static std::shared_ptr<TempDirectoryPath> files_;
  /// Map from table name to list of file system paths.
  inline static std::unordered_map<std::string, std::vector<std::string>>
      tableFilePaths_;
  inline static std::unique_ptr<melon::CPUThreadPoolExecutor> schemaExecutor_;
};

/// Reads all results from 'runner'.
std::vector<RowVectorPtr> readCursor(
    std::shared_ptr<runner::LocalRunner> runner);

} // namespace kumo::pollux::exec::test
